from fastapi import APIRouter, Header
from starlette.requests import Request

from model.db.mapper.well_mapper import get_dev_unit, get_reservoir_by_devunits, get_well_by_reservoir, get_core, \
    update_core_data, update_dev_unit, update_reservoir_unit, update_pit_unit, get_core_dataList, update_core_param, \
    update_core_basic, update_core_dataList, get_core_param, get_reservoir_list, get_well_list, get_core_by_well, \
    get_core_dataList_by_pit
from model.params.well_params import ReservoirParam, GetWellParam, CoreDataParam, UpdateDevParam, UpdateResevoirParam, \
    UpdatePitParam, UpdateCoreParam, UpdateCoreBasic, UpdateCoreDataList, WellParam
from util import success_res, obj_to_dict, SuccessCode
from util.public_func import need_serialize, objs_to_dict, transform_Datasets

well_router = APIRouter(prefix="/well", tags=['油田信息'])


@well_router.post("/dev_unit", name="获取开发单元(父子关系)")
async def get_unit(token=Header('token')):
    unit = get_dev_unit()
    # 创建一个空列表来存储下级部门
    units = objs_to_dict(unit)
    departments = {}
    for item in units:
        unit_id = item["DvpUnit_ID"]
        unit_name = item["DvpUnit_Name"]
        parent_unit_name = item["Belong_DvpUnit_Name"]
        unit_rank = int(item["DvpUnit_Rank"])
        remarks = item["Remark"]
        department = {
            "unit_id": unit_id,
            "unit_name": unit_name,
            "Belong_DvpUnit_Name": None,  # 添加一个字段来记录上级部门
            "unit_rank": unit_rank,
            "remarks": remarks,
            "children": None
        }
        departments[unit_name] = department
        if parent_unit_name:
            if parent_unit_name in departments:
                parent_department = departments[parent_unit_name]
                department["Belong_DvpUnit_Name"] = parent_department['unit_name']  # 设置当前部门的上级部门
                if parent_department["children"] is None:
                    parent_department["children"] = []
                parent_department["children"].append(department)
            else:
                print(f"无法找到上级部门: {parent_unit_name}")
    department_structure = [
        department for department in departments.values() if department["Belong_DvpUnit_Name"] is None
    ]

    return success_res(objs_to_dict(department_structure), SuccessCode.GET)


@well_router.post("/dev_unitList", name="获取开发单元")
async def get_unit(token=Header('token')):
    unit = get_dev_unit()
    return success_res(objs_to_dict(unit), SuccessCode.GET)


@well_router.post("/update_dev", name="更新开发单元")
async def update_dev(parm: UpdateDevParam, token=Header('token')):
    add = update_dev_unit(parm.add, parm.delete, parm.update)
    return success_res({'add': add}, SuccessCode.PUT)


@well_router.post("/reservoir_unit", name="根据开发单元获取储层单元")
async def get_reservoir_unit(reservoir_param: ReservoirParam, token=Header('token')):
    reservoir_unit = get_reservoir_by_devunits(reservoir_param.devUnits)
    return success_res(objs_to_dict(reservoir_unit), SuccessCode.GET)


@well_router.post("/reservoir_unit_list", name="获取储层单元列表")
async def get_reservoir_unit_list(token=Header('token')):
    reservoir_unit = get_reservoir_list()
    return success_res(objs_to_dict(reservoir_unit), SuccessCode.GET)


@well_router.post("/update_reservoir", name="更新储层单元")
async def update_reservoir(reservoir_param: UpdateResevoirParam, token=Header('token')):
    add = update_reservoir_unit(reservoir_param.add, reservoir_param.delete, reservoir_param.update)
    return success_res({'add': add}, SuccessCode.GET)


@well_router.post("/pit", name="根据储层获取井")
async def get_well_unit(well_param: GetWellParam, token=Header('token')):
    well_unit = get_well_by_reservoir(well_param.rsvUnits)
    return success_res(objs_to_dict(well_unit), SuccessCode.GET)


@well_router.post("/pitList", name="获取井列表")
async def get_well_unit(token=Header('token')):
    well_unit = get_well_list()
    return success_res(objs_to_dict(well_unit), SuccessCode.GET)


@well_router.post("/update_pit", name="更新井")
async def update_pit(param: UpdatePitParam, token=Header('token')):
    newWell = update_pit_unit(param.add, param.delete, param.update)
    return success_res({'add': objs_to_dict(newWell)}, SuccessCode.GET)


@well_router.post("/core", name="根据井号获取岩芯")
async def get_core_list(param: WellParam, token=Header('token')):
    core_list = get_core_by_well(param.wellUnits)
    return success_res(objs_to_dict(core_list), SuccessCode.GET)


@well_router.post("/core_list", name="获取岩心列表")
async def get_core_list(token=Header('token')):
    core_list = get_core()
    return success_res(objs_to_dict(core_list), SuccessCode.GET)


@well_router.post("/update_core", name="更新岩芯")
async def update_core(param: UpdateCoreBasic, token=Header('token')):
    add = update_core_basic(param.add, param.delete, param.update)
    return success_res({'add': add}, SuccessCode.GET)


@well_router.post("/get_coreParm", name="获取岩芯参数")
async def core_param(token=Header('token')):
    params = get_core_param()
    return success_res(objs_to_dict(params), SuccessCode.GET)


@well_router.post("/update_coreParm", name="更新岩芯参数")
async def update_pit(param: UpdateCoreParam, token=Header('token')):
    add = update_core_param(param.add, param.delete, param.update)
    return success_res({'add': add}, SuccessCode.GET)


@well_router.post("/core_dataList", name="获取岩心参数数据")
async def get_coreList(token=Header('token')):
    core_list = objs_to_dict(get_core_dataList())
    modify = transform_Datasets(core_list)
    return success_res(sorted(modify, key=lambda x: x["Core_ID"]), SuccessCode.GET)


@well_router.post("/core_dataList_by_pit", name="根据井获取岩心参数数据")
async def get_coreList(param: WellParam, token=Header('token')):
    core_list = objs_to_dict(get_core_dataList_by_pit(param.wellUnits))
    modify = transform_Datasets(core_list)
    return success_res(sorted(modify, key=lambda x: x["Core_ID"]), SuccessCode.GET)


@well_router.post("/update_core_data", name="更新岩心参数数据")
async def update_core_List(param: UpdateCoreDataList, token=Header('token')):
    add = update_core_dataList(param.add, param.delete, param.update)
    return success_res({'add': add}, SuccessCode.GET)

# @well_router.post("/core_data", name="更新岩心参数数据")
# async def get_core_list(param: CoreDataParam, token=Header('token')):
#     tag = update_core_data(param.Core_ID, param.Parameter_ID, param.value)
#     if tag == 1:
#         return success_res('', SuccessCode.PUT)
